package com.li.job.once;

import com.google.common.util.concurrent.ListenableFuture;
import com.li.job.properties.QWenEmbeddingProperties;
import com.li.job.properties.QdrantProperties;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.stereotype.Component;

import java.nio.file.FileSystems;
import java.nio.file.Path;
import java.nio.file.PathMatcher;
import java.nio.file.Paths;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ExecutionException;

//ApplicationRunner 和CommandLineRunner
//同: 都是在容器启动末尾执行的预加载操作
//异: 如果两者都有 ,则ApplicationRunner优先执行 ; 参数不同CommandLineRunner只允许传String参数 , ApplicationRunner接口的方法参数ApplicationArguments（是个对象）

@Component
public class RagResourceInit implements ApplicationRunner {
    @Resource
    private QWenEmbeddingProperties qWenProperties;
    @Resource
    private QdrantProperties qdrantProperties;

        @Override
        public void run(ApplicationArguments args) throws Exception {
            Path path = Paths.get(Objects.requireNonNull(getClass().getClassLoader().getResource("rag")).toURI());
            System.out.println(path);
            //pathMather 路径匹配器 ;glob:*.txt ==> 匹配当前目录下(不包括子目录) 的所有 .txt 文件
            PathMatcher pathMatcher = FileSystems.getDefault().getPathMatcher("glob:*.txt");
            // 指定要加载的目录（例如 resources/rag）
            List<Document> documents = FileSystemDocumentLoader.loadDocuments(path, pathMatcher);
            EmbeddingModel embeddingModel =  QwenEmbeddingModel.builder()
                            .modelName(qWenProperties.qwenModelName)
                                    .apiKey(qWenProperties.qwenKey)
                                    .build();

            collectionIsExist(qdrantProperties.collectionName);
            //这个不会创建collection ,如果没有connection 还得需要客户端连接
            EmbeddingStore<TextSegment> embeddingStore = QdrantEmbeddingStore.builder()
                    .host(qdrantProperties.host)
                    .port(qdrantProperties.port)
                    .collectionName(qdrantProperties.collectionName)
                    .build();
            documents.stream()
                    //flatMap 把一个流转为另一个流
                    .flatMap(document -> documentByRegexSplitter(document).stream())
                    .forEach(textSegment -> {
                        Embedding embedding =embeddingModel.embed(textSegment.text()).content();
                        embeddingStore.add(embedding,textSegment);
                      });

    }
    private List<TextSegment> documentByRegexSplitter(Document document){
        DocumentByRegexSplitter documentByRegexSplitter = new DocumentByRegexSplitter(
                "[？?。．;!；]",
                "。",
                500,
                30,
                new DocumentByCharacterSplitter(500,0));
        return documentByRegexSplitter.split(document);
    }
    private void collectionIsExist(String collectionName){
            //创建Qdrant客户端
        QdrantClient client = new QdrantClient(
                QdrantGrpcClient.newBuilder(
                        qdrantProperties.host,
                        qdrantProperties.port,
                        //false表示不器用Https/SSL加密通信 ,当使用https连接时设置为true
                        false
                ).build()
        );
        try {
            Boolean b = client.collectionExistsAsync(collectionName).get();
            if(!b){//如果不存在集合则创建集合
            client.createCollectionAsync(
                    collectionName, //集合名称
                    Collections.VectorParams.newBuilder()
                            .setDistance(Collections.Distance.Cosine)
                            .setSize(qdrantProperties.collectionSize)
                            .build()
            ).get();}
        } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }


}
